#include "cd.hpp"
#include <iostream>
#include <memory>

std::map<std::string, std::vector<std::string>> Cd::fns;
std::ofstream Cd::output_file;
std::vector<std::pair<std::string, std::string>> Cd::loop_tables;
std::set<std::string> Cd::buildin_fns;

void Cd::init(std::string input_filename, std::string output_filename) {
    std::ifstream input_file(input_filename);
    std::string content((std::istreambuf_iterator<char>(input_file)), std::istreambuf_iterator<char>());
    src = content;
    index = 0;
    output_file.open(output_filename);

    symbols = {
        "=", ",",
        "[", "]",
        "(", ")",
        "{", "}",
        ";",
    };

    op_table = {
        { "<=", 7 }, { ">=", 7 },
        { "==", 6 }, { "!=", 6 },
        { "&&", 5 },
        { "||", 4 },

        { "*", 3 }, { "/", 3 }, { "%", 3 },
        { "+", 2 }, { "-", 2 },
        { "<", 1 }, { ">", 1 },
    };

    buildin_fns = { "exit", "print", "input" };

    next_token();
}

void Cd::next_token() {
    token_val = "";
    while(src.length() > index) {
        if(src[index] == ' ' || src[index] == '\t' || src[index] == '\n' || src[index] == '\r') {
            index ++;
            continue;
        }

        if(isalpha(src[index]) || src[index] == '$' || src[index] == '_') {
            std::string ident = "";
            while(src.length() > index && (isalpha(src[index]) || isdigit(src[index]) || src[index] == '$' || src[index] == '_')) {
                ident += src[index ++];
            }

            if(ident == "while") {
                token = TOKEN_WHILE;
            } else if(ident == "for") {
                token = TOKEN_FOR;
            } else if(ident == "break") {
                token = TOKEN_BREAK;
            } else if(ident == "continue") {
                token = TOKEN_CONTINUE;
            } else if(ident == "if") {
                token = TOKEN_IF;
            } else if(ident == "else") {
                token = TOKEN_ELSE;
            } else if(ident == "fn") {
                token = TOKEN_FN;
            } else if(ident == "return") {
                token = TOKEN_RET;
            } else if(ident == "asm") {
                token = TOKEN_ASM;
            } else {
                token = TOKEN_IDENT;
                token_val = ident;
            }

            return;
        }

        if(isdigit(src[index])) {
            std::string num = "";
            while(src.length() > index && (isdigit(src[index]) || src[index] == '.')) {
                num += src[index ++];
            }

            token = TOKEN_NUM;
            token_val = num;

            return;
        }

        if(src[index] == '"') {
            index ++;   // "

            std::string str = "";
            while(src[index] != '"') {
                if(src[index] == '\\') {
                    index ++;

                    switch(src[index]) {
                        case '"': str += "\\\""; break;
                        case 'a': str += '\a'; break;
                        case 't': str += '\t'; break;
                        case 'f': str += '\f'; break;
                        case 'n': str += '\n'; break;
                        case 'r': str += '\r'; break;
                        case 'v': str += '\v'; break;
                        case '0': str += '\0'; break;
                        case '\\': str += "\\\\"; break;
                        default:  str += src[index]; break;
                    }

                    index ++;
                } else {
                    str += src[index ++];
                }
            }

            token = TOKEN_STRING;
            token_val = str;

            index ++;   // "

            return;
        }

        if(src[index] == '#') {
            while(src.length() > index && src[index] != '\n' && src[index] != '\r') {
                index ++;
            }

            index ++;
            continue;
        }

        std::string symbol1 = "";
        symbol1 += src[index ++];
        std::string symbol2 = symbol1;
        if(src.length() > index) {
            symbol2 += src[index];
        }

        if(op_table.count(symbol1) || symbols.count(symbol1)
        || op_table.count(symbol2) || symbols.count(symbol2)) {
            token = TOKEN_SYMBOL;
            if(op_table.count(symbol2) || symbols.count(symbol2)) {
                token_val = symbol2;
                index ++;
            } else {
                token_val = symbol1;
            }

            return;
        } else {
            std::cerr << "ERROR: undefined char: " << src[index] << '\n';
            exit(-1);
        }
    }

    token = TOKEN_EOF;
    token_val = "";
}

void Cd::eat_token(std::string eat) {
    next_token();
    if(token_val != eat) {
        std::cerr << "ERROR: Grammatical errors" << '\n';
        exit(-1);
    }
}

std::shared_ptr<Cd::BaseAST> Cd::parse_value() {
    if(token == TOKEN_NUM) {
        double number = std::stod(token_val);
        next_token();
        return std::make_shared<NumberAST>(number);
    } else if(token == TOKEN_IDENT) {
        std::string ident = token_val;
        next_token();

        if(token_val == "(" || token_val == "[") {
            std::shared_ptr<BaseAST> table;
            if(token_val == "(") {
                next_token();
                std::vector<std::shared_ptr<BaseAST>> args;
                while(token_val != ")") {
                    args.push_back(parse_expr());
                    
                    if(token_val == ",") {
                        next_token();
                    }
                }
                
                next_token();
                table = std::make_shared<CallAST>(ident, args);
            } else {
                table = std::make_shared<VarAST>(ident);
            }

            if(token_val == "[") {
                std::shared_ptr<GetAST> target = std::make_shared<GetAST>(table, nullptr);
                while(token_val == "[") {
                    next_token();
                    target -> key = parse_expr();
                    next_token();
                    target = std::make_shared<GetAST>(target, nullptr);
                }

                return target -> table;
            } else {
                return table;
            }
        } else {
            return std::make_shared<VarAST>(ident);
        }
    } else if(token_val == "(") {
        next_token();
        std::shared_ptr<BaseAST> target = parse_expr();
        next_token();

        return target;
    }
}

std::shared_ptr<Cd::BaseAST> Cd::parse_expr() {
    if(token == TOKEN_STRING) {
        std::string string = token_val;
        next_token();
        return std::make_shared<StringAST>(string);
    }

    if(token_val == "{") {
        next_token();   // {
        int i = 0;
        std::map<std::string, std::shared_ptr<BaseAST>> table;
        while(token_val != "}") {
            if(token == TOKEN_IDENT) {
                std::string key = token_val;
                eat_token("=");
                next_token();
                table[key] = parse_expr();
            } else {
                table[std::to_string(i ++)] = parse_expr();
            }

            if(token_val == ",") {
                next_token();
            }
        }
        next_token();

        return std::make_shared<TableAST>(table);
    }

    std::vector<std::string> ops;
    std::vector<std::shared_ptr<BaseAST>> nums;
    if(token_val == "+" || token_val == "-") {
        std::string op = token_val;
        next_token();
        nums.push_back(std::make_shared<CalcAST>(op, parse_value(), nullptr, true));
    }

    while(token != TOKEN_EOF && token_val != "}" && token_val != "]" && token_val != "," && token_val != "=" && token_val != ")" && token_val != ";") {
        if(op_table.count(token_val)) {
            if(ops.size() >= 1 && op_table[token_val] <= op_table[ops.back()]) {
                nums[nums.size() - 2] = std::make_shared<CalcAST>(ops.back(), nums.back(), nums[nums.size() - 2]);
                nums.pop_back();
                ops.pop_back();
            }
            ops.push_back(token_val);
            next_token();
        } else {
            nums.push_back(parse_value());
        }
    }

    for(int i = ops.size() - 1; i >= 0; i --) {
        nums[i] = std::make_shared<CalcAST>(ops[i], nums[i], nums[i + 1]);
    }
    
    return nums[0];
}

std::shared_ptr<Cd::BlockAST> Cd::parser() {
    std::vector<std::shared_ptr<BaseAST>> exprs;
    while(token != TOKEN_EOF && token_val != "}") {
        if(token_val == ";") {
            next_token();
            continue;
        }

        switch(token) {
            case TOKEN_WHILE: {
                next_token();
                next_token();
                std::shared_ptr<BaseAST> condition = parse_expr();
                next_token();
                next_token();
                std::shared_ptr<BaseAST> body = parser();
                next_token();

                exprs.push_back(std::make_shared<WhileAST>(condition, body));
                break;
            }

            case TOKEN_FOR: {
                next_token();
                next_token();
                std::string var_name = token_val;
                next_token();
                next_token();
                std::shared_ptr<BaseAST> start_expr = parse_expr();
                next_token();
                std::shared_ptr<BaseAST> end_expr = parse_expr();
                next_token();
                std::shared_ptr<BaseAST> step_expr = parse_expr();
                next_token();
                next_token();
                std::shared_ptr<BaseAST> body = parser();
                next_token();

                exprs.push_back(std::make_shared<ForAST>(var_name, start_expr, end_expr, step_expr, body));
                break;
            }

            case TOKEN_BREAK: {
                next_token();
                exprs.push_back(std::make_shared<BreakAST>());
                next_token();

                break;
            }

            case TOKEN_CONTINUE: {
                next_token();
                exprs.push_back(std::make_shared<ContinueAST>());
                next_token();

                break;
            }

            case TOKEN_IF: {
                next_token();
                next_token();
                std::shared_ptr<BaseAST> condition = parse_expr();
                next_token();
                next_token();
                std::shared_ptr<BaseAST> if_body = parser();
                next_token();

                std::shared_ptr<BaseAST> else_body = nullptr;
                if(token == TOKEN_ELSE) {
                    next_token();
                    next_token();
                    else_body = parser();
                    next_token();
                }

                exprs.push_back(std::make_shared<IfAST>(condition, if_body, else_body));
                break;
            }

            case TOKEN_FN: {
                next_token();
                std::string fn_name = token_val;

                eat_token("(");
                next_token();
                std::vector<std::string> arg_names;
                while(token_val != ")") {
                    arg_names.push_back(token_val);
                    next_token();

                    if(token_val == ",") {
                        next_token();
                    }
                }
                
                eat_token("{");
                next_token();
                std::shared_ptr<BaseAST> body = parser();
                next_token();

                fns[fn_name] = arg_names;
                exprs.push_back(std::make_shared<FnAST>(fn_name, arg_names, body));
                break;
            }

            case TOKEN_RET: {
                next_token();
                exprs.push_back(std::make_shared<RetAST>(parse_expr()));
                next_token();
                break;
            }

            case TOKEN_ASM: {
                next_token();
                std::string body = "";
                while(src[index] != '}') {
                    body += src[index ++];
                }
                eat_token("}");
                exprs.push_back(std::make_shared<AsmAST>(body));
                break;
            }

            default: {
                std::shared_ptr<BaseAST> target = parse_expr();

                if(token_val == "=") {
                    next_token();
                    std::shared_ptr<BaseAST> value = parse_expr();
                    next_token();

                    exprs.push_back(std::make_shared<AssignAST>(target, value));
                } else {
                    exprs.push_back(target);
                    next_token();
                }
                break;
            }
        }
    }

    return std::make_shared<BlockAST>(exprs);
}

void Cd::WhileAST::codegen() {
    std::string condition_label = std::to_string(rand());
    output_file << "!" << condition_label << '\n';
    condition -> codegen();

    std::string begin_label = std::to_string(rand());
    output_file << "jc %" << begin_label << '\n';
    std::string end_label = std::to_string(rand());
    output_file << "jmp %" << end_label << '\n';

    output_file << "!" << begin_label << '\n';
    loop_tables.push_back({ condition_label, end_label });
    body -> codegen();
    loop_tables.pop_back();
    output_file << "jmp %" << condition_label << '\n';

    output_file << "!" << end_label << '\n';
}

void Cd::ForAST::codegen() {
    start_expr -> codegen();
    output_file << "store $" << var_name << '\n';

    std::string start_label = std::to_string(rand());
    std::string end_label = std::to_string(rand());
    std::string condition_label = std::to_string(rand());
    output_file << "!" << start_label << '\n';

    loop_tables.push_back({ condition_label, end_label });
    body -> codegen();
    loop_tables.pop_back();

    output_file << "!" << condition_label << '\n';
    output_file << "load $" << var_name << '\n';

    step_expr -> codegen();
    output_file << "calc [+]\n";
    output_file << "store $" << var_name << '\n';
    output_file << "load $" << var_name << '\n';

    end_expr -> codegen();
    output_file << "calc [<=]\n";
    output_file << "jc %" << start_label << '\n';
    
    output_file << "!" << end_label << '\n';
}

void Cd::BreakAST::codegen() {
    output_file << "jmp %" << loop_tables.back().second << '\n';
}

void Cd::ContinueAST::codegen() {
    output_file << "jmp %" << loop_tables.back().first << '\n';
}

void Cd::IfAST::codegen() {
    condition -> codegen();
    std::string if_label = std::to_string(rand());
    std::string else_label = std::to_string(rand());
    std::string end_label = std::to_string(rand());

    output_file << "jc %" << if_label << '\n';
    output_file << "jmp %" << else_label << '\n';

    output_file << "!" << if_label << '\n';
    if_body -> codegen();
    output_file << "jmp %" << end_label << '\n';

    output_file << "!" << else_label << '\n';
    if(else_body) {
        else_body -> codegen();
    }

    output_file << "!" << end_label << '\n';
}

void Cd::FnAST::codegen() {
    std::string end_label = std::to_string(rand());
    output_file << "jmp %" << end_label << '\n';
    output_file << "!" << fn_name << '\n';
    body -> codegen();
    output_file << "!" << end_label << '\n';
}

void Cd::RetAST::codegen() {
    returnValue -> codegen();
    output_file << "ret\n";
}

void Cd::AsmAST::codegen() {
    output_file << body;
}

void Cd::StringAST::codegen() {
    output_file << "load \"" << string << "\"\n";
}

void Cd::NumberAST::codegen() {
    output_file << "load @" << number << '\n';
}

void Cd::VarAST::codegen() {
    if(left_mode) {
        output_file << "store $" << var_name << '\n';
    } else {
        output_file << "load $" << var_name << '\n';
    }
}

void Cd::TableAST::codegen() {
    output_file << "table [new]\n";
    for(auto item : table) {
        output_file << "load \"" << item.first << "\"\n";
        item.second -> codegen();
        output_file << "table [assign]\n";
    }
}

void Cd::CalcAST::codegen() {
    if(is_union) {
        num1 -> codegen();
        output_file << "load @" << op << "1\n";
        output_file << "calc [*]\n";
    } else {
        num1 -> codegen();
        num2 -> codegen();
        output_file << "calc [" << op << "]\n";
    }
}

void Cd::CallAST::codegen() {
    if(buildin_fns.count(fn_name)) {
        for(auto arg : args) {
            arg -> codegen();
        }

        output_file << "fn [" << fn_name << "]\n";
        return;
    }

    std::vector<std::string> arg_names = fns[fn_name];
    output_file << "table [new]\n";
    for(int i = 0; i < arg_names.size(); i ++) {
        output_file << "load \"" << arg_names[i] << "\"\n";
        args[i] -> codegen();
        output_file << "table [assign]\n";
    }
    output_file << "call %" << fn_name << '\n';
}

void Cd::AssignAST::codegen() {
    value -> codegen();

    left -> left_mode = true;    // 开启左值模式
    left -> codegen();
}

void Cd::GetAST::codegen() {
    if(left_mode) {
        std::string tmp_name = std::to_string(rand());
        output_file << "store $" << tmp_name << '\n';

        table -> codegen();
        key -> codegen();
        output_file << "load $" << tmp_name << '\n';
        output_file << "table [assign]\n";
    } else {
        table -> codegen();
        key -> codegen();
        output_file << "table [get]\n";
    }
}

void Cd::BlockAST::codegen() {
    for(auto expr : exprs) {
        expr -> codegen();
    }
}

#ifdef DEBUG
int main() {
    Cd cd;
    cd.init("test", "test.s");
    cd.parser() -> codegen();
}
#endif